"""
This file provides functions and data structures for load and writing data
"""

import numpy 

class Feature:
    """
    Class to store information about an MS1 feature. The identifiers are given by the 
    order of teh features in the input file 
    """
    
    def __init__(self, identifier, rt, mass):
        self.id = str(identifier)
        self.rt = float(rt)
        self.mass = float(mass)

    def __str__(self):
        return "%s:%g:%g" % (self.id, self.rt, self.mass)
  

class Peptide:
    """
    Class to store information about a peptide
    """

    def __init__(self, sequence, pred_rt, th_mass, proteins):
        self.sequence = sequence
        self.pred_rt = float(pred_rt)
        self.th_mass = float(th_mass)
        self.proteins = proteins
        self.matched_features = []
        self.mass_interval = (0.0, 0.0)
        
    def get_begin(self):
        return self.mass_interval[0]

    def get_end(self):
        return self.mass_interval[1]

    def __eq__(self, other):
        return self.sequence == other.sequence and self.rt == other.rt

    def __str__(self):
        return "%s:%g:%g:%s" % (self.sequence, self.pred_rt, self.th_mass, \
                ":".join(self.proteins))
 

class Protein:
    """
    Class to store information about a protein
    """
    
    def __init__(self, name, peptides, probs, matched):
        self.name = name 
        self.peptides = peptides
        self.probs = probs 
        self.matched = matched
        self.pvalue = numpy.inf
    
    def __str__(self):
        matched = len([d for d in self.matched if d == 1])
        return "%s\t%d\t%d\t%g" % (self.name, len(self.peptides), \
                matched, self.pvalue)   
    
    
    def detailed_str(self):        
        peptides = " ".join([p + ":" + str(prob) for (p, prob) in \
                zip(self.peptides, self.probs)])
        matched_peptides = " ".join([self.peptides[i] + ":" + str(self.probs[i]) for \
                i in xrange(len(self.peptides)) if self.matched[i] == 1])
        if len(matched_peptides) == 0:
            matched_peptides = "[]"
        return self.name + ":" + str(self.pvalue) + "\n" + peptides + "\n" + \
                matched_peptides + "\n"                 
        
  
def write_proteins_detailed(proteins, output_file):
    outf = open(output_file, "w")
    outf.write("# Each protein is described on three lines including:\n")
    outf.write("# Protein_name:p value\n")
    outf.write("# Peptides, as sequence:probability_null_model\n")
    outf.write("# Matched peptides, as sequence:probability_null_model" + \
               " or [] if no matched peptide \n")    
    for p in proteins:
        outf.write("%s\n" % (p.detailed_str(), ))
    outf.close()


def write_proteins_brief(proteins, output_file):
    outf = open(output_file, "w")
    outf.write("Protein\tN_peptides\tN_matched\tp_value\n")
    for p in proteins:
        outf.write("%s\n" % (str(p), ))
    outf.close()

    
def load_peptide_file(input_file): 
    """
    Input: filename including one peptide per line, the first line being header
            Each line should include the following tab-separated fields:
            Peptide_sequence Predicted_retention_time Theoretical_mass 
    Output: list of Peptide objects
    """
    in_handler = open(input_file)
    in_handler.readline()
    peptides = []
    for line in in_handler:
        fields = line.split("\t")
        proteins = [field.strip() for field in fields[3:]]
        peptides.append(Peptide(fields[0], fields[1], fields[2], proteins))
    in_handler.close()
    return peptides


def load_feature_file(input_file): 
    """
    Input: filename including one feature per line, the first line being header
            Each line should include the following tab-separated fields:
            Feature_id Observed_retention_time Observed_mass 
    Output: list of Feature objects
    """
    in_handler = open(input_file)
    in_handler.readline()
    features = []
    for line in in_handler:
        fields = line.split("\t")
        features.append(Feature(*fields))
    in_handler.close()
    return features       
 

def write_peptides(peptides, output_file, only_matched = False): 
    """
    Input: 
    peptides - list of Peptide objects to be written to the output file 
    output_file - name of the output file 
    only_matched - only the peptides with at least one matching MS1-feature will
            be included; by default all peptides are written to the output file 
    """
    out_handler = open(output_file, "w")
    out_handler.write("Peptide:predicted_rt:theoretical_mass:proteins\t" +\
            "Matching_features(id:rt:mass)\n")
    if only_matched:
        peptides = [p for p in peptides if len(p.matched_features) > 0]
    for peptide in peptides:
        out_handler.write(str(peptide))
        for feat in peptide.matched_features:
            out_handler.write("\t" + str(feat))
        out_handler.write("\n")
    out_handler.close()


def load_peptides_matched(input_file):
    """
    Input: file including peptides and the features to which they match
            One peptide per line, format: Peptide:rt:mass 
            Feature(line_in_input_file:rt:mass), all tab separated
            The first line is a header

    Output: list of Peptide objects
    """
    in_handler = open(input_file)
    in_handler.readline()
    peptides = []
    for line in in_handler:
        fields = line.split("\t")
        peptide_data = fields[0].split(":")
        proteins = [p.strip() for p in peptide_data[3:]]
        peptide = Peptide(peptide_data[0], peptide_data[1], peptide_data[2], \
                proteins)
        peptide.matched_features = [Feature(*fields[i].split(":")) \
                for i in xrange(1, len(fields))]
        peptides.append(peptide)
    in_handler.close()
    return peptides      


def write_null_probs(dict_probs, output_file):
    outf = open(output_file, "w")
    outf.write("Peptide_length\tNull_model_probability\n")
    for l, prob in dict_probs.iteritems():
        outf.write("%d\t%g\n" % (l, prob))
    outf.close()
